#!/usr/bin/env python3
import argparse, grpc, os, sys, struct, threading, time, random
from time import sleep

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../utils/')
)
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

ETH_HEADER_FORMAT = "!6s6sH"
ETH_HEADER_LENGTH = struct.calcsize(ETH_HEADER_FORMAT)


def format_ipv4_rules(p4info_helper, switch, match_data, action_data):
    if isinstance(match_data, list):
        match_data = tuple(match_data)
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": match_data},
        action_name="MyIngress.ipv4_forward",
        action_params=action_data
    )
    switch.WriteTableEntry(table_entry)


def configure_ecn_check(p4info_helper, switch, ecn_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.check_ecn",
        action_name="MyEgress.mark_ecn",
        action_params=ecn_params,
        default_action=True
    )
    switch.WriteTableEntry(table_entry)


def setup_clone_session(p4info_helper, switch, session_id, replica_list):
    session_entry = p4info_helper.buildCloneSessionEntry(
        clone_session_id=session_id,
        replicas=replica_list
    )
    switch.WritePREEntry(session_entry)


def _packetin_getter(sw):
    if hasattr(sw, "PacketIn"):
        return sw.PacketIn
    if hasattr(sw, "get_stream_packet"):
        return sw.get_stream_packet
    return None


def start_packetin_reader(sw):
    getter = _packetin_getter(sw)
    if getter is None:
        print(f"[WARN] Packet-In getter not found in p4runtime_lib for {sw.name}. Monitoring disabled for this switch.")
        return

    def _reader():
        try:
            while True:
                try:
            
                    resp = None
                    try:
                        resp = getter(timeout=2.0)
                    except TypeError:
                    
                        resp = getter()
                    if resp is None:
                        continue
                    if resp.WhichOneof("update") == "packet":
                        pkt = resp.packet.payload
                        if len(pkt) >= ETH_HEADER_LENGTH + 2:
                            ip_tos = pkt[ETH_HEADER_LENGTH + 1]
                            ecn_value = ip_tos & 0x03
                            print(f"[{sw.name}] ECN value: {ecn_value}")
                            if ecn_value == 3:
                                print(f"[{sw.name}] Congestion detected!!")
                except Exception as ie:
                
                    print(f"[WARN] PacketIn read error on {sw.name}: {ie}")
                    sleep(0.5)
        except Exception as e:
            print(f"[WARN] packet-in reader for {sw.name} stopped: {e}")

    t = threading.Thread(target=_reader, daemon=True)
    t.start()


def initialize_switch_connections():

    switches = {}
    LOGDIR = "/tmp/p4rt-logs"
    os.makedirs(LOGDIR, exist_ok=True)
    switch_configs = [
        ('s1', '127.0.0.1:50051', 0, f'{LOGDIR}/s1-p4runtime-requests.txt'),
        ('s2', '127.0.0.1:50052', 1, f'{LOGDIR}/s2-p4runtime-requests.txt'),
        ('s3', '127.0.0.1:50053', 2, f'{LOGDIR}/s3-p4runtime-requests.txt')
    ]
    for name, addr, dev_id, log_file in switch_configs:
        switches[name] = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name=name, address=addr, device_id=dev_id, proto_dump_file=log_file
        )
    return switches


def deploy_forwarding_rules(p4info_helper, switches):

    s1_rules = [
        (["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 2}),
        (["10.0.1.11", 32], {"dstAddr": "08:00:00:00:01:11", "port": 1}),
        (["10.0.2.0", 24], {"dstAddr": "08:00:00:00:02:00", "port": 3}),
        (["10.0.3.0", 24], {"dstAddr": "08:00:00:00:03:00", "port": 4})
    ]
    s2_rules = [
        (["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2}),
        (["10.0.2.22", 32], {"dstAddr": "08:00:00:00:02:22", "port": 1}),
        (["10.0.1.0", 24], {"dstAddr": "08:00:00:00:01:00", "port": 3}),
        (["10.0.3.0", 24], {"dstAddr": "08:00:00:00:03:00", "port": 4})
    ]
    s3_rules = [
        (["10.0.3.3", 32], {"dstAddr": "08:00:00:00:03:03", "port": 1}),
        (["10.0.1.0", 24], {"dstAddr": "08:00:00:00:01:00", "port": 2}),
        (["10.0.2.0", 24], {"dstAddr": "08:00:00:00:02:00", "port": 3})
    ]
    for rule in s1_rules:
        format_ipv4_rules(p4info_helper, switches['s1'], *rule)
    for rule in s2_rules:
        format_ipv4_rules(p4info_helper, switches['s2'], *rule)
    for rule in s3_rules:
        format_ipv4_rules(p4info_helper, switches['s3'], *rule)


def _make_unique_election_id():
    low = (int(time.time() * 1e9) ^ os.getpid() ^ random.getrandbits(32)) & ((1 << 64) - 1)
    return (0, low)


def acquire_mastership(switches):
    for sw in switches.values():
        eid = _make_unique_election_id()
        try:
            sw.MasterArbitrationUpdate(election_id=eid)
            print(f"[{sw.name}] MasterArbitrationUpdate OK, election_id={eid}")
        except grpc.RpcError as e:
            print(f"[{sw.name}] MasterArbitrationUpdate failed.")
            printGrpcError(e)
            raise


def configure_network_features(p4info_helper, switches, threshold, enable_clone=False):
    ecn_config = {"ecn_threshold": int(threshold)}
    for sw in switches.values():
        configure_ecn_check(p4info_helper, sw, ecn_config)

    if enable_clone:
        clone_replicas = [{"egress_port": 255, "instance": 1}]
        for sw in switches.values():
            try:
                setup_clone_session(p4info_helper, sw, 100, clone_replicas)
            except grpc.RpcError as e:
                print("WARN: clone session setup failed; continue without clone.")
                printGrpcError(e)


def monitor_network(switches):
    print("\nMonitoring network congestion ... (Ctrl-C to stop)")
    for sw in switches.values():
        start_packetin_reader(sw)
    try:
        while True:
            sleep(60)
    except KeyboardInterrupt:
        print("\nShutting down.")


def main(p4info_file_path, bmv2_file_path, monitor: bool):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    switches = initialize_switch_connections()
    acquire_mastership(switches)

    for sw in switches.values():
        sw.SetForwardingPipelineConfig(
            p4info=p4info_helper.p4info,
            bmv2_json_file_path=bmv2_file_path
        )

    try:
        threshold = int(input("Please input the threshold of the queue: ").strip())
    except Exception:
        threshold = 0
        print("Input invalid. Fallback to threshold=0")

    deploy_forwarding_rules(p4info_helper, switches)
    configure_network_features(p4info_helper, switches, threshold, enable_clone=monitor)

    if monitor:
        monitor_network(switches)
    else:
        print("\nController configured. Packet-In monitoring is disabled (use --monitor to enable).")
        try:
            while True:
                sleep(60)
        except KeyboardInterrupt:
            print("\nShutting down.")


def validate_file_paths(parser, args):
    if not os.path.exists(args.p4info):
        parser.print_help()
        print(f"\np4info file not found: {args.p4info}\nHave you run 'make'?")
        return False
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nBMv2 JSON file not found: {args.bmv2_json}\nHave you run 'make'?")
        return False
    return True


if __name__ == '__main__':
    ap = argparse.ArgumentParser(description='P4Runtime Controller')
    ap.add_argument('--p4info', default='./build/ecn.p4.p4info.txtpb', type=str,
                    help='p4info proto in text format from p4c')
    ap.add_argument('--bmv2-json', default='./build/ecn.json', type=str,
                    help='BMv2 JSON file from p4c')
    ap.add_argument('--monitor', action='store_true',
                    help='Enable Packet-In monitoring (requires clone_to_cpu in P4).')
    args = ap.parse_args()

    if validate_file_paths(ap, args):
        try:
            main(args.p4info, args.bmv2_json, monitor=args.monitor)
        except grpc.RpcError as e:
            printGrpcError(e)
        finally:
            ShutdownAllSwitchConnections()
